#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Alibaba Cloud OSS Automatic Backup Tool
- Automatically backs up specified files/directories to Alibaba Cloud OSS at scheduled times
"""

import os
import sys
import time
import logging
import datetime
import schedule
import configparser
import oss2
import hashlib
import json
from pathlib import Path

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler("oss_backup.log"),
        logging.StreamHandler(sys.stdout)
    ]
)
logger = logging.getLogger(__name__)

class OssBackupTool:
    def __init__(self, config_file="config.ini"):
        """Initialize the OSS backup tool with configuration"""
        self.config_file = config_file
        self.config = self._load_config()
        
        # Initialize OSS client
        self.auth = oss2.Auth(
            self.config.get('oss', 'access_key_id'),
            self.config.get('oss', 'access_key_secret')
        )
        self.bucket = oss2.Bucket(
            self.auth,
            self.config.get('oss', 'endpoint'),
            self.config.get('oss', 'bucket_name')
        )
        
        # Parse backup schedule
        self.backup_time = self.config.get('schedule', 'backup_time')
        
        # Parse backup sources
        self.backup_sources = self._parse_backup_sources()
        
        # Get storage limit configuration (in GB)
        self.storage_limit_gb = float(self.config.get('oss', 'storage_limit_gb', fallback='450'))
        self.storage_threshold_percent = float(self.config.get('oss', 'storage_threshold_percent', fallback='90'))
        
        # Load file hash cache
        self.hash_cache_file = "file_hash_cache.json"
        self.hash_cache = self._load_hash_cache()
        
        # Skip unchanged files setting
        self.skip_unchanged = self.config.getboolean('backup', 'skip_unchanged', fallback=True)
        
        # Show progress for large files
        self.show_progress = self.config.getboolean('backup', 'show_progress', fallback=True)
        self.large_file_threshold = int(self.config.get('backup', 'large_file_threshold_mb', fallback='10')) * 1024 * 1024  # MB to bytes
        
    def _load_config(self):
        """Load configuration from the config file"""
        if not os.path.exists(self.config_file):
            self._create_default_config()
            logger.info(f"Created default configuration file: {self.config_file}")
            logger.info("Please edit the configuration file with your OSS credentials and backup settings")
            sys.exit(0)
            
        config = configparser.ConfigParser()
        config.read(self.config_file)
        return config
    
    def _create_default_config(self):
        """Create a default configuration file"""
        config = configparser.ConfigParser()
        
        config['oss'] = {
            'access_key_id': 'YOUR_ACCESS_KEY_ID',
            'access_key_secret': 'YOUR_ACCESS_KEY_SECRET',
            'endpoint': 'http://oss-cn-hangzhou.aliyuncs.com',
            'bucket_name': 'your-bucket-name',
            'storage_limit_gb': '450',  # 90% of 500GB
            'storage_threshold_percent': '90'  # When to start cleaning up
        }
        
        config['schedule'] = {
            'backup_time': '02:00',  # Default backup time at 2 AM
        }
        
        config['backup'] = {
            'sources': r'C:\path\to\backup1, D:\path\to\backup2',
            'exclude_patterns': '*.tmp, *.log, .git/',
            'retention_days': '30',  # Keep backups for 30 days
            'skip_unchanged': 'true',  # Skip files that haven't changed since last backup
            'show_progress': 'true',  # Show progress bar for large files
            'large_file_threshold_mb': '10'  # Files larger than this will show progress (in MB)
        }
        
        with open(self.config_file, 'w') as f:
            config.write(f)
    
    def _load_hash_cache(self):
        """Load the file hash cache from disk"""
        if os.path.exists(self.hash_cache_file):
            try:
                with open(self.hash_cache_file, 'r') as f:
                    return json.load(f)
            except Exception as e:
                logger.warning(f"Failed to load hash cache: {str(e)}")
                return {}
        return {}
    
    def _save_hash_cache(self):
        """Save the file hash cache to disk"""
        try:
            with open(self.hash_cache_file, 'w') as f:
                json.dump(self.hash_cache, f)
        except Exception as e:
            logger.warning(f"Failed to save hash cache: {str(e)}")
    
    def _calculate_file_hash(self, file_path):
        """Calculate MD5 hash of a file"""
        hash_md5 = hashlib.md5()
        file_size = os.path.getsize(file_path)
        show_progress = self.show_progress and file_size > self.large_file_threshold
        
        try:
            with open(file_path, "rb") as f:
                if show_progress:
                    # 显示哈希计算进度条
                    sys.stdout.write(f"\rCalculating hash for {os.path.basename(file_path)}: 0% (0/{file_size/1024/1024:.2f} MB)")
                    sys.stdout.flush()
                    
                    # 使用更大的缓冲区以提高性能
                    buffer_size = 8192 * 1024  # 8MB buffer
                    bytes_read = 0
                    
                    while True:
                        data = f.read(buffer_size)
                        if not data:
                            break
                        hash_md5.update(data)
                        bytes_read += len(data)
                        
                        # 更新进度条
                        percent = int(100 * bytes_read / file_size)
                        sys.stdout.write(f"\rCalculating hash for {os.path.basename(file_path)}: {percent}% ({bytes_read/1024/1024:.2f}/{file_size/1024/1024:.2f} MB)")
                        sys.stdout.flush()
                    
                    # 完成时打印换行
                    sys.stdout.write("\n")
                    sys.stdout.flush()
                else:
                    # 对于小文件，一次性读取全部内容
                    for chunk in iter(lambda: f.read(4096), b""):
                        hash_md5.update(chunk)
                        
            return hash_md5.hexdigest()
        except Exception as e:
            logger.warning(f"Failed to calculate hash for {file_path}: {str(e)}")
            return None
            
    def _get_file_metadata(self, file_path):
        """Get file metadata including size and modification time"""
        try:
            stat = os.stat(file_path)
            return {
                "size": stat.st_size,
                "mtime": stat.st_mtime
            }
        except Exception as e:
            logger.warning(f"Failed to get metadata for {file_path}: {str(e)}")
            return None
    
    def _parse_backup_sources(self):
        """Parse backup sources from config"""
        sources_str = self.config.get('backup', 'sources')
        return [src.strip() for src in sources_str.split(',')]
    
    def _get_object_key(self, local_path):
        """Generate an OSS object key for the given local path"""
        timestamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
        filename = os.path.basename(local_path)
        return f"backups/{timestamp}/{filename}"
    
    def _progress_callback(self, bytes_transferred, total_bytes):
        """Callback for upload progress"""
        if total_bytes:
            rate = int(100 * (float(bytes_transferred) / float(total_bytes)))
            sys.stdout.write(f"\rProgress: {rate}% ({bytes_transferred}/{total_bytes} bytes)")
            sys.stdout.flush()
            
            # Print newline when complete
            if bytes_transferred == total_bytes:
                sys.stdout.write("\n")
                sys.stdout.flush()
    
    def get_bucket_usage(self):
        """Get the current bucket usage in bytes"""
        total_size = 0
        object_count = 0
        
        # Iterate through all objects to calculate total size
        for obj in oss2.ObjectIterator(self.bucket, prefix="backups/"):
            total_size += obj.size
            object_count += 1
            
        return {
            "total_size_bytes": total_size,
            "total_size_gb": total_size / (1024 * 1024 * 1024),
            "object_count": object_count
        }
    
    def check_storage_space(self):
        """Check if storage space is approaching the limit"""
        usage = self.get_bucket_usage()
        storage_limit_bytes = self.storage_limit_gb * 1024 * 1024 * 1024
        usage_percent = (usage["total_size_bytes"] / storage_limit_bytes) * 100
        
        logger.info(f"Current OSS storage usage: {usage['total_size_gb']:.2f} GB ({usage_percent:.2f}% of {self.storage_limit_gb} GB limit)")
        
        if usage_percent >= self.storage_threshold_percent:
            logger.warning(f"Storage usage exceeds threshold of {self.storage_threshold_percent}%, will clean up old backups")
            return True
        return False
    
    def get_backup_folders_sorted_by_date(self):
        """Get all backup folders sorted by date (oldest first)"""
        folders = {}
        
        # Extract unique folder timestamps
        for obj in oss2.ObjectIterator(self.bucket, prefix="backups/"):
            parts = obj.key.split('/')
            if len(parts) >= 2:
                folder_timestamp = parts[1]
                if folder_timestamp not in folders:
                    try:
                        # Convert timestamp to datetime for sorting
                        timestamp = datetime.datetime.strptime(folder_timestamp, "%Y%m%d-%H%M%S")
                        folders[folder_timestamp] = timestamp
                    except ValueError:
                        # Skip if format doesn't match
                        continue
        
        # Sort folders by timestamp (oldest first)
        sorted_folders = sorted(folders.items(), key=lambda x: x[1])
        return [folder[0] for folder in sorted_folders]
    
    def delete_backup_folder(self, folder_timestamp):
        """Delete all objects in a specific backup folder"""
        deleted_size = 0
        object_count = 0
        
        prefix = f"backups/{folder_timestamp}/"
        for obj in oss2.ObjectIterator(self.bucket, prefix=prefix):
            deleted_size += obj.size
            self.bucket.delete_object(obj.key)
            object_count += 1
            
        logger.info(f"Deleted backup folder {folder_timestamp} - {object_count} files, {deleted_size / (1024 * 1024):.2f} MB freed")
        return deleted_size
    
    def clean_up_space_if_needed(self, needed_bytes=0):
        """Clean up space if needed, deleting the oldest backups first"""
        # 检查存储空间是否超过阈值
        if not self.check_storage_space() and needed_bytes == 0:
            return False
            
        # 获取当前存储使用情况
        usage = self.get_bucket_usage()
        storage_limit_bytes = self.storage_limit_gb * 1024 * 1024 * 1024
        
        # 计算可用空间
        available_space = storage_limit_bytes - usage["total_size_bytes"]
        
        # 只有当实际需要空间超过可用空间时才清理
        if needed_bytes > available_space:
            # 计算需要释放的空间量（至少10%的存储限制或所需空间）
            min_space_to_free = max(storage_limit_bytes * 0.1, needed_bytes - available_space)
            
            logger.info(f"Need to free at least {min_space_to_free / (1024 * 1024 * 1024):.2f} GB of space")
            
            # 获取按日期排序的备份文件夹
            folders = self.get_backup_folders_sorted_by_date()
            if not folders:
                logger.warning("No backup folders found to clean up")
                return False
                
            # 删除最旧的备份直到释放足够空间
            freed_space = 0
            for folder in folders:
                freed_space += self.delete_backup_folder(folder)
                if freed_space >= min_space_to_free:
                    logger.info(f"Freed up enough space: {freed_space / (1024 * 1024 * 1024):.2f} GB")
                    return True
                    
            logger.warning(f"Could only free {freed_space / (1024 * 1024 * 1024):.2f} GB, which may not be enough")
            return freed_space > 0
        
        # 空间充足，不需要清理
        return False
    
    def has_file_changed(self, file_path):
        """Check if a file has changed since the last backup. Returns (changed, new_hash)"""
        if not self.skip_unchanged:
            return (True, None)  # We don't calculate hash now, but we will in backup? Actually, we don't know the new_hash, so we return None and let backup calculate it.

        # Get current file metadata
        metadata = self._get_file_metadata(file_path)
        if not metadata:
            return (True, None)  # Assume changed, and we don't have a new_hash (it will be calculated during backup)

        normalized_path = os.path.normpath(file_path)

        # If the file is in the cache, check metadata
        if normalized_path in self.hash_cache:
            cached_data = self.hash_cache[normalized_path]

            # If metadata matches, file hasn't changed
            if cached_data["size"] == metadata["size"] and cached_data["mtime"] == metadata["mtime"]:
                return (False, None)

            # Metadata changed, so we need to check the hash
            # Calculate the new hash
            file_size = metadata["size"]
            if self.show_progress and file_size > self.large_file_threshold:
                logger.info(f"File {os.path.basename(file_path)} has metadata changed, calculating hash...")

            new_hash = self._calculate_file_hash(file_path)
            if new_hash is None:
                return (True, None)  # Assume changed

            # Compare with cached hash
            if cached_data["hash"] == new_hash:
                # File content hasn't changed, so we return unchanged
                # But note: we don't update the cache here (even metadata) because we are not backing up
                return (False, None)
            else:
                # Content has changed
                return (True, new_hash)
        else:
            # File not in cache -> new file, so we need to back it up
            # Calculate the hash now to avoid recalculating during backup?
            file_size = metadata["size"]
            if self.show_progress and file_size > self.large_file_threshold:
                logger.info(f"New file {os.path.basename(file_path)}, calculating hash...")

            new_hash = self._calculate_file_hash(file_path)
            if new_hash is None:
                return (True, None)  # Assume changed, but we don't have a hash? Then during backup we will recalc?
            return (True, new_hash)
    
    def backup_file(self, file_path):
        """Upload a single file to OSS"""
        if not os.path.exists(file_path):
            logger.error(f"File not found: {file_path}")
            return False
        
        try:
            # Check if file has changed since last backup
            changed, new_hash = self.has_file_changed(file_path)
            if not changed:
                logger.info(f"Skipping unchanged file: {file_path}")
                return True
                
            # Get file size and check if we need to clean up space
            file_size = os.path.getsize(file_path)
            self.clean_up_space_if_needed(file_size)
            
            object_key = self._get_object_key(file_path)
            logger.info(f"Backing up {file_path} to OSS key: {object_key}")
            
            # For large files, use progress reporting
            show_progress = self.show_progress and file_size > self.large_file_threshold
            
            # Upload the file to OSS
            if show_progress:
                # Define progress callback function for current upload
                # OSS SDK可能使用不同版本的回调格式，这里适配两种格式
                def _progress_callback(consumed_bytes, total_bytes=None):
                    if total_bytes is None:
                        # 旧版SDK格式: (consumed_bytes, total_bytes)
                        return
                        
                    rate = int(100 * (float(consumed_bytes) / float(total_bytes)))
                    sys.stdout.write(f"\rUploading {os.path.basename(file_path)}: {rate}% ({consumed_bytes/1024/1024:.2f}/{total_bytes/1024/1024:.2f} MB)")
                    sys.stdout.flush()
                    
                    # Print newline when complete
                    if consumed_bytes == total_bytes:
                        sys.stdout.write("\n")
                        sys.stdout.flush()
                
                # Use the progress callback
                result = self.bucket.put_object_from_file(
                    object_key, 
                    file_path,
                    progress_callback=_progress_callback
                )
            else:
                result = self.bucket.put_object_from_file(object_key, file_path)
            
            if result.status == 200:
                if show_progress:
                    logger.info(f"Successfully backed up {file_path} ({file_size/1024/1024:.2f} MB)")
                else:
                    logger.info(f"Successfully backed up {file_path}")
                
                # Update hash cache AFTER successful backup
                normalized_path = os.path.normpath(file_path)
                metadata = self._get_file_metadata(file_path)
                if metadata:
                    self.hash_cache[normalized_path] = {
                        "hash": new_hash if 'new_hash' in locals() else self._calculate_file_hash(file_path),
                        "size": metadata["size"],
                        "mtime": metadata["mtime"],
                        "last_backup": datetime.datetime.now().isoformat()
                    }
                return True
            else:
                logger.error(f"Failed to backup {file_path}, status: {result.status}")
                return False
                
        except Exception as e:
            logger.error(f"Error backing up {file_path}: {str(e)}")
            return False
    
    def backup_directory(self, directory_path):
        """Backup all files in a directory recursively"""
        if not os.path.exists(directory_path):
            logger.error(f"Directory not found: {directory_path}")
            return False
        
        exclude_patterns = [p.strip() for p in self.config.get('backup', 'exclude_patterns', fallback='').split(',')]
        success_count = 0
        failure_count = 0
        skipped_count = 0
        
        # Calculate total size first to check if we need to clean up space
        total_size = 0
        files_to_backup = []
        
        for root, dirs, files in os.walk(directory_path):
            # Apply exclusion patterns to directories
            dirs[:] = [d for d in dirs if not any(Path(os.path.join(root, d)).match(pattern) for pattern in exclude_patterns if pattern)]
            
            for file in files:
                file_path = os.path.join(root, file)
                
                # Skip files matching exclude patterns
                if any(Path(file_path).match(pattern) for pattern in exclude_patterns if pattern):
                    logger.debug(f"Skipping excluded file: {file_path}")
                    continue
                    
                # Check if file has changed
                if self.skip_unchanged:
                    changed, _ = self.has_file_changed(file_path)   # We ignore the new_hash because we are not backing up this file (if unchanged) and if we do back up, backup_file will handle the cache
                    if not changed:
                        logger.info(f"Skipping unchanged file: {file_path}")
                        skipped_count += 1
                        continue
                    
                files_to_backup.append(file_path)
                total_size += os.path.getsize(file_path)
        
        # Clean up space if needed before starting backup
        if total_size > 0:
            self.clean_up_space_if_needed(total_size)
        
        # Now perform the backup
        for file_path in files_to_backup:
            if self.backup_file(file_path):
                success_count += 1
            else:
                failure_count += 1
        
        logger.info(f"Directory backup summary for {directory_path}: {success_count} files succeeded, {failure_count} files failed, {skipped_count} files skipped (unchanged)")
        
        # Save hash cache after backup
        self._save_hash_cache()
        
        return failure_count == 0
    
    def run_backup(self):
        """Run the backup process for all configured sources"""
        logger.info("Starting backup process...")
        
        timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        logger.info(f"Backup started at: {timestamp}")
        
        # Check storage space before starting backup
        self.check_storage_space()
        
        success = True
        for source in self.backup_sources:
            source = source.strip()
            if not source:
                continue
                
            logger.info(f"Processing backup source: {source}")
            
            if os.path.isfile(source):
                if not self.backup_file(source):
                    success = False
            elif os.path.isdir(source):
                if not self.backup_directory(source):
                    success = False
            else:
                logger.error(f"Backup source not found: {source}")
                success = False
        
        # Save hash cache after all backups
        self._save_hash_cache()
        
        end_timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        if success:
            logger.info(f"Backup completed successfully at: {end_timestamp}")
        else:
            logger.warning(f"Backup completed with errors at: {end_timestamp}")
        
        return success
    
    def cleanup_old_backups(self):
        """Delete backups older than retention_days"""
        retention_days = int(self.config.get('backup', 'retention_days', fallback='30'))
        if retention_days <= 0:
            logger.info("Backup retention is set to keep all backups")
            return
            
        logger.info(f"Cleaning up backups older than {retention_days} days")
        
        cutoff_time = datetime.datetime.now() - datetime.timedelta(days=retention_days)
        cutoff_str = cutoff_time.strftime("%Y%m%d")
        
        # List all backup objects
        for obj in oss2.ObjectIterator(self.bucket, prefix="backups/"):
            try:
                # Extract date from the object key (format: backups/YYYYMMDD-HHMMSS/filename)
                parts = obj.key.split('/')
                if len(parts) >= 2:
                    date_str = parts[1].split('-')[0]  # Get YYYYMMDD part
                    if date_str < cutoff_str:
                        logger.info(f"Deleting old backup: {obj.key}")
                        self.bucket.delete_object(obj.key)
            except Exception as e:
                logger.error(f"Error while cleaning up old backup {obj.key}: {str(e)}")
    
    def schedule_backup(self):
        """Schedule regular backups"""
        logger.info(f"Scheduling daily backup at {self.backup_time}")
        schedule.every().day.at(self.backup_time).do(self.run_backup)
        schedule.every().day.at(self.backup_time).do(self.cleanup_old_backups)
        
        while True:
            schedule.run_pending()
            time.sleep(60)  # Check every minute

def main():
    """Main entry point for the backup tool"""
    try:
        backup_tool = OssBackupTool()
        
        # If arguments are provided, perform immediate backup
        if len(sys.argv) > 1 and sys.argv[1] == "--now":
            backup_tool.run_backup()
        else:
            # Otherwise schedule regular backups
            logger.info("OSS Backup Tool started")
            backup_tool.schedule_backup()
            
    except KeyboardInterrupt:
        logger.info("Backup tool stopped by user")
    except Exception as e:
        logger.error(f"Unhandled exception: {str(e)}")
        return 1
        
    return 0

if __name__ == "__main__":
    sys.exit(main()) 